module Zeno.Term where

import Prelude ()
import StdImports
import qualified Data.Set as Set

data Term a
  = Var !a
  | App !(Term a) !(Term a)
  | Err
  deriving ( Eq, Ord, Functor, Foldable, Traversable )
  
instance NFData a => NFData (Term a) where
  rnf Err = ()
  rnf (Var a) = rnf a
  rnf (App t1 t2) = rnf t1 `seq` rnf t2
  
instance WithinTraversable (Term a) (Term a) where
  mapWithinM f (App x y) =
    f =<< App `liftM` mapWithinM f x `ap` mapWithinM f y
  mapWithinM f t = f t

type TermSubstitution a = Substitution (Term a)

equalTerms :: Eq a => Term a -> Term a -> Bool
equalTerms (Var x) (Var y) = x == y
equalTerms (App x1 y1) (App x2 y2) = x1 == x2 && y1 == y2
equalTerms _ _ = False

isVar :: Term a -> Bool
isVar (Var {}) = True
isVar _ = False

fromVar :: Term a -> a
fromVar (Var a) = a

isApp :: Term a -> Bool
isApp (App {}) = True
isApp _ = False

collapseErrors :: Term a -> Term a
collapseErrors (App Err _) = Err
collapseErrors (App x y) = App (collapseErrors x) (collapseErrors y)
collapseErrors t = t

instance Show a => Show (Term a) where
  show Err = "_|_"
  show (Var a) = show a
  show (App (App (Var x) y) z) 
    | not (isAlpha $ head $ show x) =
        showBracketed y ++ " " ++ show x ++ " " ++ showBracketed z
  show (App x y) = show x ++ " " ++ showBracketed y
  
showBracketed :: Show a => Term a -> String
showBracketed t@(App x y) = "(" ++ show t ++ ")"
showBracketed t = show t
 
termFunction :: Term a -> a
termFunction (Var a) = a
termFunction (App t1 t2) = termFunction t1

flattenTerm :: Term a -> [Term a]
flattenTerm (App f a) = flattenTerm f ++ [a]
flattenTerm t = [t]

unflattenTerm :: [Term a] -> Term a
unflattenTerm [] = Err
unflattenTerm ts = foldl1 App ts

splitTerms :: (Term a -> [Term a]) -> Term a -> [Term a]
splitTerms f t@(App t1 t2) =
  let t1s = splitTerms f t1
      t2s = splitTerms f t2
  in f t ++ (map (App t1) t2s) ++ (map (flip App t2) t1s)
splitTerms f t = f t
  
splitTermsM :: Monad m => (Term a -> m [Term a]) -> Term a -> m [Term a]
splitTermsM f t@(App t1 t2) = do
  t1s <- splitTermsM f t1
  t2s <- splitTermsM f t2
  ft <- f t
  return $ ft ++ (map (App t1) t2s) ++ (map (flip App t2) t1s)
splitTermsM f t = f t

splitMapTerms :: Functor f => 
  (Term a -> [f (Term a)]) -> Term a -> [f (Term a)]
splitMapTerms f t@(App t1 t2) = 
  let t1s = splitMapTerms f t1
      t2s = splitMapTerms f t2
  in f t ++ (map (fmap (App t1)) t2s) ++ (map (fmap (flip App t2)) t1s) 
splitMapTerms f t = f t 

splitMapTermsM :: (Functor f, Monad m) =>
  (Term a -> m [f (Term a)]) -> Term a -> m [f (Term a)]
splitMapTermsM f t@(App t1 t2) = do
  t1s <- splitMapTermsM f t1
  t2s <- splitMapTermsM f t2
  t' <- f t
  return $ t' ++ (map (fmap (App t1)) t2s) 
    ++ (map (fmap (flip App t2)) t1s)
splitMapTermsM f t = f t
  
removeFactors :: Eq a => [Term a] -> [Term a]
removeFactors terms = filter (not . isFactor) terms
  where isFactor term = any (flip containsStrictly term) terms
   
instance Ord a => Unifiable (Term a) where
  unify t1 t2 =
    let (t1_func : t1_args) = flattenTerm t1
        (t2_func : t2_args) = flattenTerm t2
    in if (t1_func == t2_func) && (length t1_args == length t2_args)
         then mconcat (zipWith unify t1_args t2_args)
         else Set.singleton (t1, t2)
      

